package org.fatewa.engine.magic;


import cn.hutool.core.collection.ConcurrentHashSet;
import cn.hutool.core.util.TypeUtil;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import jakarta.servlet.http.HttpServletRequest;
import lombok.Data;
import lombok.SneakyThrows;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.fatewa.engine.design.magic.Subject;
import org.jetbrains.annotations.NotNull;
import org.springframework.context.ApplicationContext;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Type;
import java.sql.Connection;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 数据拦截器, 可动态进行 SQL 修改，实现数据拦截
 *
 * @author fatewa
 */
@Data
@Accessors(fluent = true, chain = true)
@Slf4j
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class RayquazaSqlInterceptor implements Interceptor {

    /**
     * mybatis plus 分页 auto count 语句
     */
    private static final String MP_PAGE = "_mpCount";

    /**
     * 静态存储类型对应的主体拦截器
     * {@link AbstractSubjectInterceptor}
     * 一个类型可以对应多个主体拦截器
     */
    private final Map<Class<?>, Collection<AbstractSubjectInterceptor<?, ?>>> SUBJECTS = new HashMap<>();

    /**
     * Spring 上下文对象
     */
    private final ApplicationContext context;

    /**
     * Base Mapper 的方法缓存
     */
    private final static Set<String> BASE_MAPPER_METHODS = new HashSet<>();


    public final static Set<String> WHITELIST = new HashSet<>();
    private final static String SUBJECTS_METHOD = "subjects";
    private final static String SKIP_METHOD = "skip";

    static {
        Arrays.stream(BaseMapper.class.getDeclaredMethods()).forEach((it) -> BASE_MAPPER_METHODS.add(it.getName()));
    }

    /**
     * 初始化拦截器映射关系，便于后续处理
     */
    public void doCacheIfNotExist() {
        // 已经处理过则跳过
        if (SUBJECTS.size() > 0) {
            return;
        }
        //noinspection rawtypes
        Map<String, AbstractSubjectInterceptor> interceptorMap = context.getBeansOfType(AbstractSubjectInterceptor.class);
        //noinspection rawtypes
        for (AbstractSubjectInterceptor it : interceptorMap.values()) {
            // 获取所有作用域
            List<AbstractSubjectInterceptor.ActionScope> scopes = ((AbstractSubjectInterceptor<?, ?>) it).actionScopes();

            // 遍历所有的作用域
            for (AbstractSubjectInterceptor.ActionScope scope : scopes) {
                // 目标实体
                Class<?> clazz = scope.clazz();
                // 使用 Set 进行去重
                SUBJECTS.putIfAbsent(clazz, new ConcurrentHashSet<>());
                Collection<AbstractSubjectInterceptor<?, ?>> interceptors = SUBJECTS.get(clazz);
                // 添加拦截器
                interceptors.add(it);
            }
        }
    }

    public RayquazaSqlInterceptor(ApplicationContext context) {
        this.context = context;
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        HttpServletRequest request = currentRequest();
        StackTraceElement[] stackTrace = Thread.currentThread().getStackTrace();
        for (StackTraceElement stackTraceElement : stackTrace) {
            String className = stackTraceElement.getClassName();
            // 白名单存在时直接返回，减少类反射开销
            if (WHITELIST.contains(className)) {
                return invocation.proceed();
            }
            // 遇到 subjects 方法时，不进行拦截
            String methodName = stackTraceElement.getMethodName();
            if (SUBJECTS_METHOD.equals(methodName) || SKIP_METHOD.equals(methodName)) {
                Class<?> originClass = Class.forName(className);
                if (AbstractSubjectInterceptor.class.isAssignableFrom(originClass)) {
                    WHITELIST.add(className);
                    return invocation.proceed();
                }
            }
        }

        // 非 Web 上下文, 不进行拦截处理
        if (null == request) {
            return invocation.proceed();
        }

        // 扫描拦截器，并进行缓存
        doCacheIfNotExist();

        StatementHandler statementHandler = PluginUtils.realTarget(invocation.getTarget());
        MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");

        final String id = mappedStatement.getId();

        // mapper 类
        String mapper = id.substring(0, id.lastIndexOf("."));

        // 调用的 mapper 方法
        String method = id.substring(id.lastIndexOf(".") + 1);

        // 遇到分页插件的前置 count，去除后缀获得原始方法
        if (method.endsWith(MP_PAGE)) {
            method = method.replaceAll(MP_PAGE, "");
        }

        // 仅拦截 BaseMapper 中的方法
        if (!BASE_MAPPER_METHODS.contains(method)) {
            return invocation.proceed();
        }

        // 获取 mapper 类型
        final Class<?> mapperClass = getMybatisMapperClass(mapper);

        // 仅拦截 mybatis plus 的 base mapper
        if (null != mapperClass && !BaseMapper.class.isAssignableFrom(mapperClass)) {
            return invocation.proceed();
        }

        // 根据 Mapper 类型的泛型获取目标实体类型
        Class<?> entity = getEntityClassByMapper(mapperClass);

        // 获取实体类的拦截器
        Collection<AbstractSubjectInterceptor<?, ?>> interceptors = SUBJECTS.get(entity);

        if (null != interceptors) {
            log.info("{} interceptors matched, start processing sql", entity.getName());
            String rawSql = this.parseOriginalSql(statementHandler);
            long start = System.currentTimeMillis();
            // 执行处理
            this.handleSubjectInterceptor(entity, mappedStatement, interceptors, metaObject, rawSql);
            log.info("sql used {}ms", System.currentTimeMillis() - start);
        }

        return invocation.proceed();
    }

    /**
     * 根据 Mapper {@link BaseMapper} 获取实体类型
     *
     * @param mapperClass mapper 的 class
     * @return 实体 class
     * @throws ClassNotFoundException 可能存在类型找不到的情况
     */
    @NotNull
    private static Class<?> getEntityClassByMapper(Class<?> mapperClass) throws ClassNotFoundException {
        Type argument = TypeUtil.getTypeArgument(mapperClass);
        return Class.forName(argument.getTypeName());
    }

    /**
     * 获取当前的请求
     *
     * @return 当前的请求
     */
    public HttpServletRequest currentRequest() {
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (null != attributes) {
            return attributes.getRequest();
        }
        return null;
    }

    /**
     * 执行处理主体拦截器
     *
     * @param entity          当前正在处理的实体类型
     * @param mappedStatement {@link MappedStatement}
     * @param interceptors    拦截器集合
     * @param metaObject      用于修改 sql 的元数据
     * @param rawSql          原始 sql
     */
    @SneakyThrows
    private <E, T extends Subject<E>> void
    handleSubjectInterceptor(Class<?> entity, MappedStatement mappedStatement,
                             Collection<AbstractSubjectInterceptor<?, ?>> interceptors,
                             MetaObject metaObject, String rawSql) {

        for (AbstractSubjectInterceptor<?, ?> it : interceptors) {
            // 转换为泛型
            //noinspection unchecked
            AbstractSubjectInterceptor<E, T> interceptor = (AbstractSubjectInterceptor<E, T>) it;

            // 如判断可跳过
            if (interceptor.skip()) {
                continue;
            }

            // 取出当前拥有的主体
            List<T> subjects = interceptor.subjects();
            // 取出所有的作用域定义
            List<AbstractSubjectInterceptor.ActionScope> scopes = interceptor.actionScopes();
            // 定义实例化
            List<AbstractSubjectInterceptor.ActionScopeSession> sessions =
                    scopes.stream()
                            .filter(scope -> entity.equals(scope.clazz()))
                            .map(actionScope -> {
                                AbstractSubjectInterceptor.ActionScopeSession session =
                                        new AbstractSubjectInterceptor.ActionScopeSession();
                                // 作用域类型
                                session.clazz(actionScope.clazz());
                                // 作用域字段名称
                                session.field(actionScope.field());
                                // 设置原始 sql
                                session.rawSql(rawSql);
                                // 设置调整 sql 的元数据
                                session.meta(metaObject);
                                return session;
                            }).collect(Collectors.toList());

            for (AbstractSubjectInterceptor.ActionScopeSession session : sessions) {
                String sql = interceptor.getSql(session, subjects, mappedStatement);
                // 修改 SQL
                session.sql(sql);
            }

        }
    }


    /**
     * 获取原始 SQL
     *
     * @param statementHandler Mybatis statementHandler
     * @return 原始 SQL
     */
    private String parseOriginalSql(StatementHandler statementHandler) {
        return statementHandler.getBoundSql().getSql();
    }


    /**
     * 获取 Mapper Class 实例
     *
     * @param buf 类路径
     * @return Mapper Class 实例
     */
    private Class<?> getMybatisMapperClass(String buf) {
        final String clazzName = buf.replaceAll("\\.$", "");
        try {
            return Class.forName(clazzName);
        } catch (ClassNotFoundException e) {
            return null;
        }
    }

}
